# Best-first search algorithm
# Explorative implementation.
# (c) Martijn Pieters 2004, 2005

from pqueue import PQueue

#
# Generic best-first search algorithm.
#
class SearchState:
    """Base class for best-first searches.

    Provides the interface for the search, including the closed status, the
    heuristic and the framework support for storing the state in the required
    structures (hashing and such).

    """
    closed = 0    # Not a seperate list, the states themselves maintain this
    depth = 0     # Distance from starting state
    parent = None # Reference to state that generated this

    #
    # Abstract methods to implement in concrete classes
    #
    def getInternalState(self):
        """Get the internal state of this search node.

        Should return a immutable python type that identifies this node
        uniquely.

        """
        raise NotImplementedError

    def calculateHeuristicScore(self):
        """Calculate the heuristic distance to the goal state as an int"""
        raise NotImplementedError

    def createChildren(self):
        """Create child nodes.

        Implementation can be a generator.

        """
        raise NotImplementedError

    def getGoalState(self):
        """Return the goal state"""
        raise NotImplementedError

    #
    # State implementation
    #
    def listChildren(self):
        """Returns a list of child nodes"""
        for child in self.createChildren():
            child.parent = self
            child.depth = self.depth + 1
            yield child

    _distance = None
    def getHeuristicScore(self):
        """Get the heursitic score. Value is cached."""
        if self._distance is None:
            self._distance = self.calculateHeuristicScore()
        return self._distance

    def __cmp__(self, other):
        return cmp(self.getInternalState(), other.getInternalState())

    def __hash__(self):
        return hash(self.getInternalState())


# A* search algorithm
def searchBestFirst(start):
    goal = start.getGoalState()
    open = PQueue()
    open.insert(start.getHeuristicScore(), start)
    seen = {start: start}

    while open:
        # Find lowest score and remove from open
        priority, current = open.pop()

        # Give an indication of the progress so far:
        # print 79 * '\r',
        # print "Searching, depth = %2d, f() = %2d" % (current.depth, priority),

        if current == goal:
            print "\nCreated %d states, %d states left on open." % (
                len(seen), len(open))
            return make_path(current, seen)

        # The following calculates f(n) = h(n) + g(n) for all child states
        children = [(c.getHeuristicScore() + c.depth, c)
                    for c in current.listChildren()]
        for priority, child in children:
            if not seen.has_key(child):
                seen[child] = child
                open.insert(priority, child)
            else:
                # We've seen this state before, is it a shorter path?
                oldchild = seen[child]
                if oldchild.depth > child.depth:
                    if not oldchild.closed:
                        # The longer path should be removed
                        del open[oldchild]
                    open.insert(priority, child)
                    # Make sure the shorter path is recorded.
                    seen[child] = child

        seen[current].closed = 1

    print "\nSearched all %d states, through %d levels." % (
        len(seen), max(map(lambda s: s.depth, seen.keys())))
    return None

# Use psyco if available (python JIT compiler) for the search algorithm.
# Note though, that psyco trades memory for speed, so the memory requirements
# will go up notably for large state spaces.
try:
    import psyco
    psyco.bind(searchBestFirst)
except ImportError:
    pass


def make_path(end, seen):
    """Back-track from end back to goal

    Using the parent references, reconstruct the path again. The seen structure
    holds the nodes with the shortest paths.

    """
    path = [end]
    while seen[path[-1]].parent:
        path += [seen[path[-1]].parent]
    path.reverse()
    return path


#
# Maze implementation of search algorithm
#
class MazeState(SearchState):
    """Maze location representation in a state-space search
   
    State is represented as a an (x, y) tuple.

    In order to accomodate the search algorithm, the distance from the starting
    point is tracked, a heuristic function is provided, and it's closed state
    is tracked (wether or not it already has been considered).
   
    The state includes a reference to the parent state in the graph, so the
    path through the graph can be reconstructed.
   
    """
   
    def __init__(self, maze, pos=(0, 0)):
        """Create a new maze state of at position pos.
       
        """
        self._img = maze
        self._state = pos

    def getInternalState(self):
        return self._state

    def calculateHeuristicScore(self):
        """Return the minimum distance from the goal """
        goalx, goaly = self.getGoalState()._state
        x, y = self._state
        return abs(goalx - x) + abs(goaly - y)
   
    def createChildren(self):
        """Create child states.
       
        Returns all possible child nodes as an iterable.
       
        """
        x, y = self._state
        img = self._img
        for i in (-1, 0, 1):
            for j in (-1, 0, 1):
                if (i and not j) or (j and not i):
                    pos = (x + i, y + j)
                    if pos[0] < 0 or pos[1] < 0:
                        continue
                    if pos[0] >= img.size[0] or pos[1] >= img.size[1]:
                        continue
                    r, g, b, a = self._img.getpixel(pos)
                    if not g and not b:
                        yield self.__class__(self._img, pos)
   
    def getGoalState(self):
        return self.__class__(self._img, (1, 640))


import Image
maze = Image.open('maze.png')
   
start = MazeState(maze, (639, 0))
path = searchBestFirst(start)
   
if path:
    print 'Path length:', len(path)
    path = [n.getInternalState() for n in path]
    data = ""
    for pos in path[1::2]:
        r, g, b, a = maze.getpixel(pos)
        data += chr(r)
        open('mazedata.zip', 'wb').write(data)
